#!/usr/bin/env python3
# Copyright (c) Megvii, Inc. and its affiliates.

import sys
sys.path.insert(0,"../../")

import argparse
import os

import cv2
import numpy as np

import onnxruntime

from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis

input_shape = (640,640)
image_path = 'test.jpg'
model = "yolox_s.onnx"
score_thr = 0.3
output_dir = 'outdir'
    
if __name__ == '__main__':
    print("input_shape:{}".format(input_shape))
    origin_img = cv2.imread(image_path)
    img, ratio = preprocess(origin_img, input_shape)

    session = onnxruntime.InferenceSession(model)
    print("model:{}, score_thr:{}".format(model,score_thr))

    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
    output = session.run(None, ort_inputs)
    predictions = demo_postprocess(output[0], input_shape)[0]

    boxes = predictions[:, :4]
    scores = predictions[:, 4:5] * predictions[:, 5:]

    boxes_xyxy = np.ones_like(boxes)
    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
    boxes_xyxy /= ratio
    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
    if dets is not None:
        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
        origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                         conf=score_thr, class_names=COCO_CLASSES)

    mkdir(output_dir)
    output_path = os.path.join(output_dir, os.path.basename(image_path))
    cv2.imwrite(output_path, origin_img)
